package com.liuqi.openai.chain;

import com.liuqi.openai.chain.result.ChainResult;
import com.liuqi.openai.chain.result.Content;
import com.liuqi.openai.chain.result.Function;
import com.liuqi.openai.chain.result.Think;
import com.liuqi.openai.data.message.AiMessage;
import com.liuqi.openai.data.message.ChatMessage;
import com.liuqi.openai.data.message.ToolResultMessage;
import com.liuqi.openai.data.message.UserMessage;
import com.liuqi.openai.memory.ChatMemory;
import com.liuqi.openai.model.OpenAiChatModel;
import com.liuqi.openai.model.OpenAiModelException;
import com.liuqi.openai.model.handler.StreamingResponseHandler;
import com.liuqi.openai.model.output.Response;
import com.liuqi.openai.tool.CustomTool;
import com.liuqi.openai.tool.ToolCallRequest;
import com.liuqi.openai.tool.ToolExecuteException;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 对话链的执行器
 *
 * @author liuqi
 * @date 2025/7/20
 **/
public class ChainExecutor {
    private final static Logger logger = LoggerFactory.getLogger(ChainExecutor.class);

    private static final DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");

    /**
     * 对话模型
     */
    protected final OpenAiChatModel chatModel;

    /**
     * 对话记忆
     */
    private final ChatMemory chatMemory;

    /**
     * key: tool name
     * value: CustomTool
     */
    private final Map<String, CustomTool> customToolMap;

    /**
     * 当前对话的变量数据, 对话结束后, 需要清空
     */
    private final Map<String, String> currentChatVariables;

    /**
     * 流式响应处理器
     */
    private final ChainStreamingHandler handler;

    /**
     * 这是一个临时的缓存消息的队列，在发起请求的时候，需要将临时队列数据合并到主缓存中{@link #chatMemory}
     * 每当一轮对话结束后，需要合并到主缓存中{@link #chatMemory}
     */
    private final List<ChatMessage> tempMessages = new LinkedList<>();

    /**
     * 流式响应: 用于标记思考时间耗时
     */
    private final ThinkDuration thinkDuration = new ThinkDuration();

    /**
     * 当前对话的交互消息结果集, 包含如下消息
     *
     * AI_response      AI响应内容消息
     * AI_tool_call     AI响应工具消息
     * AI_call_result   AI工具调用结果
     */
    private final ChainResult chainResult = new ChainResult();

    private UserMessage currentDatetime;

    ChainExecutor(OpenAiChatModel chatModel,
                  ChatMemory chatMemory,
                  Map<String, CustomTool> customToolMap,
                  Map<String, String> currentChatVariables,
                  final ChainStreamingHandler streamingHandler) {
        this.chatModel = chatModel;
        this.chatMemory = chatMemory;
        this.customToolMap = customToolMap;
        this.currentChatVariables = currentChatVariables;

        if (streamingHandler != null) {
            this.handler = new ChainStreamingHandler() {
                @Override
                public void onContent(Content content) {
                    try {
                        streamingHandler.onContent(content);
                    } catch (Throwable t) { }
                }
                @Override
                public void onThink(Think think) {
                    try {
                        streamingHandler.onThink(think);
                    } catch (Throwable t) {
                    }
                }
                @Override
                public void onToolCall(Function function) {
                    try {
                        streamingHandler.onToolCall(function);
                    } catch (Throwable t) {
                    }
                }
                @Override
                public void onComplete() {
                    try {
                        streamingHandler.onComplete();
                    } catch (Throwable t) {
                    }
                }
                @Override
                public void onError(Throwable error) {
                    try {
                        streamingHandler.onError(error);
                    } catch (Throwable t) {
                    }
                }
            };
        } else {
            this.handler = null;
        }

        // 调用链结果集 设置开始时间
        this.chainResult.begin();
    }

    /**
     * 对话发起执行
     *
     * @param userMessage   用户消息
     * @return
     */
    public ChainResult execute(UserMessage userMessage) throws Throwable {
        // 调用链结果集 设置用户消息
        chainResult.setUserMessage(userMessage);

        // 缓存用户消息
        addMessage(userMessage);

        try {
            // 发起请求
            execute0();

            // 调用链结果集 设置结束时间
            chainResult.end();

            // 临时的消息缓存队列合并到主缓存上
            tempMessages.stream()
                    .filter(message -> message != currentDatetime)
                    .forEach(chatMemory::add);

            // 完成
            if (handler != null) {
                handler.onComplete();
            }
        } finally {
            // 清空临时缓存消息
            tempMessages.clear();
        }

        return chainResult;
    }

    /**
     * 缓冲队列添加当前时间
     */
    public void addCurrentDatetime() {
        currentDatetime = UserMessage.from(
                "current_data_time: " + LocalDateTime.now().format(formatter)
        );
        tempMessages.add(currentDatetime);
    }

    /**
     * 添加缓存消息到临时缓存队列
     */
    public void addMessage(ChatMessage message) {
        tempMessages.add(message);
    }

    /**
     * 对话发起执行
     */
    private void execute0() throws OpenAiModelException {
        // 是否流式响应
        boolean stream = handler != null;

        if (logger.isDebugEnabled()) {
            logger.debug("chat model stream is " + stream);
        }

        // 发送消息的开始时间
        long begin = System.currentTimeMillis();

        Response<AiMessage> response;
        // 流式响应
        if (stream) {
            response = chatModel.generate(chatMessages(), chainStreamingResponseHandler(handler));
        } else {
            response = chatModel.generate(chatMessages());
        }
        AiMessage aiMessage = response.content();

        // 处理完消息的结束时间
        long end = System.currentTimeMillis();

        // 对话链结果集Ai响应消息
        chainResult.addAssistantMessage(begin, end, aiMessage, thinkDuration.getDurationAndReset(), response.tokenUsage());

        // 缓存AI消息
        addMessage(aiMessage);

        // 检查是否需要执行ToolCall
        checkToolCallAndExecute(aiMessage);
    }

    /**
     * 获取聊天消息, 主缓存消息队列和临时缓存队列进行合并
     */
    private List<ChatMessage> chatMessages() {
        return Stream.concat(chatMemory.messages().stream(), tempMessages.stream())
                .collect(Collectors.toList());
    }

    /**
     * 检查是否需要执行ToolCall
     * 如果需要执行toolcall, 执行toolcall后, 拿toolcall执行结果再次和AI进行交互, 由AI生成结果并返回
     */
    private void checkToolCallAndExecute(AiMessage aiMessage) throws OpenAiModelException {
        if (aiMessage.toolCallRequests() == null) {
            // 不进行进行toolcall
            return;
        }

        // 需要执行tool call
        for (ToolCallRequest toolCallRequest : aiMessage.toolCallRequests()) {
            String name = toolCallRequest.getName();
            String arguments = toolCallRequest.getArguments();

            if (logger.isDebugEnabled()) {
                logger.debug("tool_call execute: tool_call_name" + name + ", arguments " + arguments);
            }

            /**
             * 回调工具调用
             */
            if (handler != null) {
                handler.onToolCall(new Function(name, arguments));
            }

            // 工具执行结果
            String result;
            // 开始时间
            long begin = System.currentTimeMillis();

            // 通过名称获取 CustomTool
            CustomTool customTool = customToolMap.get(name);

            if (customTool == null) {
                // 一般不会出现这种情况
                // AI都是基于提供的 tool_calls 列表进行返回 tool_call_name
                logger.error("tool_call not found, tool_call_name = " + name);

                result = "tool_call not found, tool_call_name = " + name;
            } else {
                try {
                    // 执行外部工具调用
                    result = customTool.execute(arguments, currentChatVariables).toString();

                    if (logger.isDebugEnabled()) {
                        logger.debug("tool_call execute, tool_call_name: {}, arguments: {}, result: {}",
                                name, arguments, result);
                    }
                } catch (ToolExecuteException e) {
                    logger.error("tool call failed with, exception message: " + e.getMessage(), e);
                    result = "tool call failed with, exception message: " + e.getMessage();
                }
            }

            // 结束时间
            long end = System.currentTimeMillis();

            // 封装工具执行结果消息
            ToolResultMessage toolResultMessage = ToolResultMessage.from(name, toolCallRequest.getId(), result);

            // 对话链结果集工具执行结果消息
            chainResult.addToolResultMessage(begin, end, toolResultMessage);

            // 缓存工具执行结果消息
            addMessage(toolResultMessage);
        }

        // 拿tool_calls执行结果再次发起AI请求
        execute0();
    }

    void stop() {

    }

    /**
     * 流式响应: 用于标记思考时间耗时
     */
    private class ThinkDuration {
        Long startTime;
        Long currentTime;
        long mark() {
            long currentTime = System.currentTimeMillis();
            if (this.startTime == null) {
                this.startTime = currentTime;
            }
            this.currentTime = currentTime;
            return currentTime;
        }
        void reset() {
            this.startTime = null;
            this.currentTime = null;
        }
        Float getDurationAndReset() {
            if (currentTime == null) {
                return null;
            } else {
                float duration = (currentTime - startTime);
                reset();
                return duration;
            }
        }
    }

    /**
     * 流式响应处理
     */
    private StreamingResponseHandler chainStreamingResponseHandler(final ChainStreamingHandler handler) {
        return new StreamingResponseHandler() {
            @Override
            public void onNext(String content) {
                //logger.info("onNext: " + content);
                handler.onContent(new Content(content));
            }
            @Override
            public void onThink(String think) {
                //logger.info("onThink: " + think);
                long time = thinkDuration.mark();
                handler.onThink(new Think(String.valueOf(time), think));
            }
            @Override
            public void onComplete() {
                //logger.info("onComplete");
            }
            @Override
            public void onError(Throwable error) {
                //logger.info("onError: ", error);
                handler.onError(error);
            }
        };
    }
}
